#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Alternative approach to get WeChat miniapp URL using ADB directly
This avoids chromedriver issues by using ADB to extract information

Usage:
    python get_wechat_url_adb.py [device_id]

Example:
    python get_wechat_url_adb.py 6990f488
"""
import sys
import subprocess
import time
import re
import json
import os

def run_adb_command(command, device_id=None):
    """Run an ADB command and return the output"""
    cmd = ["adb"]
    if device_id:
        cmd.extend(["-s", device_id])
    cmd.extend(command.split())
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return result.stdout.strip()
    except subprocess.CalledProcessError as e:
        print(f"Error running ADB command: {e}")
        print(f"Error output: {e.stderr}")
        return None

def get_package_name(app_name="微信", device_id=None):
    """Get the package name for an app based on its display name"""
    packages = run_adb_command("shell pm list packages -f", device_id)
    if not packages:
        return None
    
    # For WeChat, we know it's com.tencent.mm
    if app_name == "微信":
        return "com.tencent.mm"
    
    # For other apps, we'd need to parse the output and match
    return None

def get_wechat_current_activity(device_id=None):
    """Get the current activity of WeChat"""
    output = run_adb_command("shell dumpsys activity activities | grep -E 'mResumedActivity|mFocusedActivity'", device_id)
    if not output:
        return None
    
    # Look for WeChat package name in the output
    for line in output.splitlines():
        if "com.tencent.mm" in line:
            # Extract activity name using regex
            match = re.search(r'com\.tencent\.mm/([^ ]+)', line)
            if match:
                return match.group(1)
    
    return None

def extract_debuggable_webviews(device_id=None):
    """Get a list of debuggable WebViews"""
    webviews = run_adb_command("shell cat /proc/net/unix | grep webview_devtools", device_id)
    if not webviews:
        return []
    
    # Parse the output and extract WebView information
    webview_list = []
    for line in webviews.splitlines():
        if "@webview_devtools" in line:
            # Extract the socket name and pid
            parts = line.strip().split()
            if len(parts) >= 8:
                socket = parts[7]
                pid_match = re.search(r'@webview_devtools_remote_(\d+)', socket)
                if pid_match:
                    pid = pid_match.group(1)
                    process = run_adb_command(f"shell ps -p {pid} | grep -v PID", device_id)
                    if process and "com.tencent.mm" in process:
                        webview_list.append({
                            "pid": pid,
                            "socket": socket,
                            "process": process.strip()
                        })
    
    return webview_list

def forward_webview_socket(pid, device_id=None):
    """Forward the WebView socket to a local port"""
    local_port = 9222 + int(pid) % 1000  # Generate a somewhat unique port
    result = run_adb_command(f"forward tcp:{local_port} localabstract:webview_devtools_remote_{pid}", device_id)
    if result is not None:  # Command was successful
        print(f"[*] Forwarded WebView socket for PID {pid} to port {local_port}")
        return local_port
    return None

def get_webview_pages(port):
    """Get the list of pages from a WebView debugger"""
    import requests
    try:
        response = requests.get(f"http://localhost:{port}/json/list", timeout=5)
        if response.status_code == 200:
            return response.json()
        else:
            print(f"[!] Error getting WebView pages: HTTP {response.status_code}")
    except Exception as e:
        print(f"[!] Exception getting WebView pages: {e}")
    
    return []

def execute_js_in_webview(port, page_id, js_code):
    """Execute JavaScript in a WebView page and return the result"""
    import requests
    try:
        payload = {
            "id": 1,
            "method": "Runtime.evaluate",
            "params": {
                "expression": js_code,
                "returnByValue": True
            }
        }
        
        # Find the correct WebSocket URL
        pages = get_webview_pages(port)
        ws_url = None
        
        for page in pages:
            if page.get("id") == page_id:
                ws_url = page.get("webSocketDebuggerUrl")
                break
                
        if not ws_url:
            print(f"[!] Could not find WebSocket URL for page {page_id}")
            return None
            
        # Replace the WebSocket protocol with HTTP
        http_url = ws_url.replace("ws://", "http://")
        
        # Send the request
        response = requests.post(http_url, json=payload, timeout=5)
        if response.status_code == 200:
            result = response.json()
            if "result" in result and "result" in result["result"]:
                return result["result"]["result"].get("value")
        
        print(f"[!] Error executing JavaScript: {response.text}")
    except Exception as e:
        print(f"[!] Exception executing JavaScript: {e}")
    
    return None

def extract_path_from_title(title):
    """Extract miniapp path information from a page title"""
    # Pattern for titles like: wxcc909565f9ddc6f9:wei_idiom/pages/index/index.html:VISIBLE
    wx_page_match = re.match(r'(wx[a-z0-9]+):(.+?)(?:\.html)?(?::[A-Z]+)?$', title)
    
    if wx_page_match:
        appid = wx_page_match.group(1)
        path = wx_page_match.group(2)
        
        # Check if path contains html extension
        if not path.endswith('.html'):
            path = f"{path}.html"
            
        return {
            'appid': appid,
            'path': path,
            'full_path': path
        }
    
    return None

def get_weixin_url_info(device_id=None):
    """Main function to get WeChat URL information"""
    # Step 1: Check if device is connected
    devices = run_adb_command("devices")
    if not devices or "device" not in devices:
        print("[!] No devices connected or device not authorized")
        return
    
    # Step 2: Check if WeChat is running
    activity = get_wechat_current_activity(device_id)
    if not activity:
        print("[!] WeChat is not running or not in foreground")
        print("[*] Please open WeChat and navigate to a miniapp")
        return
    
    print(f"[*] Current WeChat activity: {activity}")
    
    # Step 3: Check if we're in a miniapp
    if "appbrand" not in activity.lower() and "launcher" in activity.lower():
        print("[!] Not currently in a miniapp")
        print("[*] Please navigate to a miniapp and try again")
        return
    
    # Step 4: Get WebView debugging information
    webviews = extract_debuggable_webviews(device_id)
    
    if not webviews:
        # If no webviews found, WebView debugging might not be enabled
        print("[!] No debuggable WebViews found")
        print("[*] Please make sure WebView debugging is enabled in Developer Options")
        print("[*] Enable: Settings > Developer options > Web View debug > ON")
        return
    
    print(f"[*] Found {len(webviews)} debuggable WebViews for WeChat")
    
    # Step 5: For each WebView, forward and inspect
    results = []
    processed_pids = set()  # Track which PIDs we've already processed
    
    for webview in webviews:
        pid = webview['pid']
        
        # Skip if we've already processed this PID
        if pid in processed_pids:
            continue
            
        processed_pids.add(pid)
        
        print(f"\n[*] Inspecting WebView with PID {pid}")
        
        # Forward the WebView socket
        port = forward_webview_socket(pid, device_id)
        if not port:
            continue
        
        # Get the list of pages
        pages = get_webview_pages(port)
        if not pages:
            print(f"[!] No pages found for WebView with PID {pid}")
            continue
        
        print(f"[*] Found {len(pages)} pages in WebView with PID {pid}")
        
        # Try to find miniapp-related pages
        miniapp_pages = []
        
        for i, page in enumerate(pages):
            title = page.get("title", "")
            url = page.get("url", "")
            page_id = page.get("id", "")
            
            # Look for miniapp-related pages - check title, not just URL
            is_miniapp = False
            path_info = extract_path_from_title(title)
            
            if (path_info or 
                "appbrand" in url.lower() or 
                "appbrand" in title.lower() or
                "miniprogram" in url.lower() or 
                "miniprogram" in title.lower() or 
                "小程序" in title):
                is_miniapp = True
                
                # Extract path info from title if available
                if path_info:
                    print(f"        Appid: {path_info['appid']}")
                    print(f"        Path: {path_info['path']}")
            
            print(f"    Page {i}: {title}")
            print(f"        URL: {url}")
            print(f"        ID: {page_id}")
            print(f"        Miniapp: {'Yes' if is_miniapp else 'No'}")
            
            if is_miniapp:
                # Add path info to the page object if available
                if path_info:
                    page['path_info'] = path_info
                miniapp_pages.append(page)
        
        # For each miniapp page, try to extract URL information
        for page in miniapp_pages:
            page_id = page.get("id", "")
            print(f"\n[*] Extracting URL info from page {page_id}")
            
            # Try different JavaScript approaches
            js_tests = [
                ("window.__route__", "execute_script('return window.__route__')"),
                ("window.__queryString__", "execute_script('return window.__queryString__')"),
                ("Current URL", "execute_script('return window.location.href')"),
                ("wx object", "execute_script('return typeof wx !== \"undefined\"')"),
                ("Route data", "execute_script('var data = {}; if(typeof window.__route__ !== \"undefined\") { data.route = window.__route__; data.query = window.__queryString__; } return data;')"),
            ]
            
            page_results = {"page_id": page_id, "tests": {}}
            
            for name, code in js_tests:
                result = execute_js_in_webview(port, page_id, code.replace("execute_script('", "").replace("')", ""))
                page_results["tests"][name] = result
                print(f"    {name}: {result}")
            
            # If we have both route and query, construct the full path
            if (page_results["tests"].get("Route data") and 
                isinstance(page_results["tests"]["Route data"], dict) and
                "route" in page_results["tests"]["Route data"]):
                
                route_data = page_results["tests"]["Route data"]
                route = route_data.get("route", "")
                query = route_data.get("query", "")
                
                if route:
                    full_path = f"{route}.html{('?' + query) if query else ''}"
                    page_results["full_path"] = full_path
                    print(f"\n[+] Reconstructed path: {full_path}")
            
            results.append(page_results)
    
    # Print a summary
    print("\n[*] Summary of results:")
    if results:
        for i, result in enumerate(results):
            print(f"    Result {i+1}:")
            
            # Check for path info in the page itself
            page = next((p for p in miniapp_pages if p.get('id') == result['page_id']), None)
            if page and 'path_info' in page:
                print(f"        From title: {page['path_info']['full_path']}")
                print(f"        Appid: {page['path_info']['appid']}")
            
            # Check for full_path from JavaScript results
            if "full_path" in result:
                print(f"        From JS: {result['full_path']}")
            elif "tests" in result and "window.__route__" in result["tests"] and result["tests"]["window.__route__"]:
                route = result["tests"]["window.__route__"]
                query = result["tests"].get("window.__queryString__", "")
                full_path = f"{route}.html{('?' + query) if query else ''}"
                print(f"        From JS: {full_path}")
            else:
                # If we didn't get path from JS but have it from title, don't say "could not determine"
                if not (page and 'path_info' in page):
                    print(f"        Could not determine path from JavaScript")
            
            # Show test results that might be relevant
            if "tests" in result:
                for test_name, test_result in result["tests"].items():
                    if test_result is not None and test_result != "":
                        print(f"        {test_name}: {test_result}")
    else:
        print("    No results found")
        
        # If we found miniapp pages but couldn't extract JS results,
        # at least show the path information from the titles
        if miniapp_pages:
            print("\n    Pages found but couldn't extract JavaScript data:")
            for i, page in enumerate(miniapp_pages):
                print(f"    Page {i+1}: {page.get('title', '')}")
                if 'path_info' in page:
                    print(f"        Path: {page['path_info']['path']}")
                    print(f"        Appid: {page['path_info']['appid']}")
                    
            print("\n    Try this code in your get_wechat_url method:")
            print("""
    def get_wechat_url(self) -> str:
        # First try the JavaScript approach
        try:
            url = self.driver.execute_script('return window.__route__+".html" + (window.__queryString__ ? "?"+window.__queryString__ : "");')
            if url and url != "null.html" and url != "undefined.html":
                return url
        except Exception as e:
            pass
            
        # Fall back to extracting from window title
        try:
            title = self.driver.title
            match = re.match(r'(wx[a-z0-9]+):(.+?)(?:\.html)?(?::[A-Z]+)?$', title)
            if match:
                path = match.group(2)
                if not path.endswith('.html'):
                    path = f"{path}.html"
                return path
        except Exception as e:
            pass
            
        # Default fallback
        return "pages/index/index.html"
            """)
    
    # Cleanup - handle errors gracefully
    for webview in webviews:
        port = 9222 + int(webview['pid']) % 1000
        try:
            run_adb_command(f"forward --remove tcp:{port}", device_id)
        except Exception as e:
            print(f"[!] Error removing port forward for {port}: {e}")
    
    print("\n[*] Completed URL extraction")

if __name__ == "__main__":
    device_id = None
    if len(sys.argv) > 1:
        device_id = sys.argv[1]
    
    get_weixin_url_info(device_id) 